from ultralytics.models import YOLO
import mlflow
import waveletai
import torch
from waveletai.constants import FeatureType
import os
from mlflow.utils.file_utils import TempDir
import shutil
import click
import sys
from pathlib import Path
FILE = Path(__file__).resolve()
ROOT =FILE.parents[0]  #YOL0v5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  #add ROOT tO PATH
ROOT =Path(os.path.relpath(ROOT,Path.cwd()))#relative
#os.environ["WAVELETAI_API_URL"]= "http://192.168.2.89/api
@click.command()
@click.option("--epochs","-ep",type=int, default=1)
@click.option("--batch_size","-bs",type=int, default=64)
@click.option("--weights_file","-wf",type=str, default='yolov8n.pt')
@click.option("--training_data","-td",type=str, default="",help="数据集id")
def enter(epochs,batch_size,weights_file, training_data):
    pwd =os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, pwd)
    os.chdir(pwd)
# with mlflow.start run(nested=True):
    waveletai.init()
    if weights_file.endswith(".pt"):
        weights_file = weights_file
    else:
        pretrain_model= weights_file
        if(pretrain_model):
            print("pretrain model processing.........")
            #waveletai.set_model("1a7cdc805ad641c0854911b788ea8247")
            waveletai.download_model_version_asset(pretrain_model, 'artifacts/best.pt',"./pretrain_model")
            assert os.path.exists("pretrain_model/best.pt"), "best.pt not exists"
            weights_file ='pretrain model/best.pt'
    waveletai.download_dataset_artifacts(training_data, './dataset', unzip=True, feature_type=FeatureType.YOLOV5.value)

    # 数据处理
    with TempDir() as tmp:
        print("dataset processing..........")
        tmp_path = tmp.path()
        output_directory = tmp_path + '/'
        print("output directory:", output_directory)
        waveletai.download_dataset_artifacts(training_data, output_directory + '/dataset', unzip = True,
                                             feature_type = FeatureType.YOLOV5.value)
        raw_train = output_directory + 'dataset/train'
        raw_valid = output_directory + 'dataset/valid'
        if not os.path.isdir(output_directory + "data/train/images"):
            os.makedirs(output_directory + "data/train/images")
        if not os.path.isdir(output_directory + "data/train/labels"):
            os.makedirs(output_directory + "data/train/labels")
        if not os.path.isdir(output_directory + "data/valid/images"):
            os.makedirs(output_directory + "data/valid/images")
        if not os.path.isdir(output_directory + "data/valid/labels"):
            os.makedirs(output_directory + "data/valid/labels")
        if not os.path.isdir(output_directory + "data/test/images"):
            os.makedirs(output_directory + "data/test/images")
        if not os.path.isdir(output_directory + "data/test/labels"):
            os.makedirs(output_directory + "data/test/labels")

        files = os.listdir(output_directory + 'dataset')
        print(files)

        list_train = os.listdir(raw_train)
        print("list train:", len(list_train))
        for i in range(len(list_train)):
            if list_train[i][-4:] == '.jpg':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/train/images/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)
            elif list_train[i][-4:] == '.txt':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/train/labels/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)
        print('完成 train -复制到images/复制到labels')

        if not os.path.isdir(raw_valid):
            list_valid = os.listdir(raw_train)
            raw_valid = raw_train
        else:
            list_valid = os.listdir(raw_valid)
        print("list_valid:", len(list_valid))
        for i in range(len(list_valid)):
            if list_valid[i][-4:] == '.jpg':
                old_name = raw_valid + '/' + str(list_valid[i])
                new_name = output_directory + 'data/valid/images/' + str(list_valid[i])
                shutil.copyfile(old_name, new_name)
            elif list_valid[i][-4:] == '.txt':
                old_name = raw_valid + '/' + str(list_valid[i])
                new_name = output_directory + 'data/valid/labels/' + str(list_valid[i])
                shutil.copyfile(old_name, new_name)
        for i in range(len(list_train)):
            if list_train[i][-4:] == '.jpg':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/valid/images/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)
            elif list_train[i][-4:] == '.txt':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/valid/labels/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)

        print('完成 valid -复制到images/复制到labels')  # 频据处理结柬
        print("创建对象")
        # 银存参数
        waveletai.log_param("batch_size_value", batch_size)
        print('log params')
        waveletai.log_param("epochs_value", epochs)
        # 训练模型
        print("!!!!!!!!!!!weights_file为" + weights_file)
        model = YOLO(weights_file)
        device = 0 if torch.cuda.is_available() else "cpu"
        print('begin to train')
        model.train(data=output_directory + '/dataset/data.yaml', epochs=epochs, imgsz=640, device=device, project=str(pwd))
        print('end train')

if __name__ == "__main__":
    enter()
